import torch
from torchvision import transforms
from PIL import Image
import os

# 模型定义（你自己的结构）
from MILmodel import COPDNet

# 加载模型
model = COPDNet()
model.load_state_dict(torch.load("fold_4_model_resnet34_val_acc_1.0000_val_auc_1.0000_train_acc_0.9542_train_auc_0.9888_epoch_5.pth", map_location="cpu"))
model.eval()

# 标签
with open("labels.txt") as f:
    LABELS = [line.strip() for line in f]

# 图像预处理
#preprocess = transforms.Compose([
#    transforms.Resize((224, 224)),
#    transforms.ToTensor(),
#])

# def predict(image: Image.Image, topk=3):
#    x = preprocess(image).unsqueeze(0)
#    with torch.no_grad():
#        output = model(x)
#        probs = torch.softmax(output[0], dim=0)
#    top_probs, top_ids = torch.topk(probs, topk)
#    return {
#        "label": LABELS[top_ids[0]],
#        "confidence": float(top_probs[0]),
#        "top_3": [
#            {"label": LABELS[i], "confidence": float(p)}
#            for i, p in zip(top_ids, top_probs)
#        ]
#    }
